#include "tutorialAnalysis.hxx"
#include "FiducialVolumeDefinition.hxx"
#include "Parameters.hxx"
#include "tutorialSelection.hxx"
#include "tutorialBranchesSelection.hxx"
#include "tutorialCorrection.hxx"
#include "tutorialWeightSystematics.hxx"
#include "tutorialVariationSystematics.hxx"
#include "CategoriesUtils.hxx"
#include "BasicUtils.hxx"


/*
  A highland2 Analysis inherits in ultimate instance from AnalysisAlgorithm. 
  In this particular case an intermediate class baseAnalysis is used, 
  which does all the work that is common for tracker analysis 
  ( this package will be renamed soon to something like baseTrackerAnalysis). 
  Then tutorialAnalysis inherits from baseAnalysis, and baseAnalysis inherits from AnalysisAlgorithm.
  
  The class that does the loops (over events, over toys, etc) is AnalysisLoop (under highlandTools). 
  There you can see the analysis flow, which is as follows

  LOOP over Spills{
    InitializeSpill
    LOOP over Bunches in the Spill{
      InitializeBunch
      LOOP over Configurations{
        InitializeConfiguration
        LOOP over Toy experiments for each configuration}
	  InitializeToy
          Process (Apply Selections)
	  FinalizeToy
          FillMicroTrees (Fill toy-dependent variables in micro-trees)
	}
        FillToyVarsInMicroTrees (Fill toy-independent variables)
        FillCategories (Fill categories for color code drawing)
	FinalizeConfiguration
      }
      FinalizeBunch
    }
    FillTruthTree (Fill the truth tree)
    FinalizeSpill
  }

  The Initialize.. and Finalize... methods can be implemented by the user to do more complicated analyses, but are not mandatory


  These is the list of mandatory methods to configure the analysis (call at initialization level):
  
  - DefineSelections:     Add to the SelectionManager the selections  we have defined in other files
  - DefineSystematics:    Add to the SystematicManager the systematics to be run (defined in other files)
  - DefineCorrections:    Add to the CorrectionManager the corrections to be run (defined in other files)
  - DefineConfigurations: Add to the ConfigurationManager the configurations to be considered
  - DefineMicroTrees:     Define the micro-tree variables
  - DefineTruthTree:      Define the variables of the truth tree

  These are the mandatory methods that are run for each Toy Experiment 

  - Process:                 Run the selections
  - FillToyVarsInMicroTrees: Fill the micro-trees toy-dependent   variables (run for each toy experiment)

  These are the mandatory methods that are run for Event (Bunch)

  - FillMicroTrees:          Fill the micro-trees toy-independent variables 
  - FillCategories:          Fill the micro-tree variables used for color code drawing

  These are the mandatory methods that are run for each Spill

  - FillTruthTree:           Fill the truth tree variables
  - CheckFillTruthTree:      Check whether to include or not  a given true vertex in the truth tree
  - CheckFillRooTrackerTree: Check whether to include or not  a given true vertex in the RooTracker tree     
*/


//********************************************************************
tutorialAnalysis::tutorialAnalysis(AnalysisAlgorithm* ana) : baseAnalysis(ana) {
//********************************************************************

  // Minimum accum level to save event into the output Micro-trees
  SetMinAccumCutLevelToSave(ND::params().GetParameterI("tutorialAnalysis.MinAccumLevelToSave"));

  // Minimum accum level to save event into the output RooTrackerTree
  SetMinAccumCutLevelToSaveRooTracker(ND::params().GetParameterI("tutorialAnalysis.MinAccumLevelToSaveRooTracker"));
  
  // Add the package version (to be stored in the "config" tree)
  ND::versioning().AddPackage("tutorialAnalysis", anaUtils::GetSoftwareVersionFromPath((std::string)getenv("TUTORIALANALYSISROOT")));
}

//********************************************************************
void tutorialAnalysis::DefineSelections(){
//********************************************************************
  
  /* In this method the user will add to the SelectionManager (accessed by  sel() ) the selections to be run, 
     defined in other files. In general an analysis has a single selection, which could have multiple branches. 
     Each of the branches is in fact a different selection (different cuts and actions), but defining them as branches 
     we ussualy gain in speed and simplicity since the steps that are common are applied only once. 
     Sometimes the user does not want to expend time merging two selections in a single one (as branches), but preffers to run 
     both independently. This is in general done for quick checks (are the selections mutually exclusive ?, etc). 
     This is possible in highland2. An example on how to do that is shown below. 
   */

  // Add two selections to the SelectionManager provided:
  // - Name of the selection (kTrackerTutorial, kTrackerTutorialBranches)
  // - Title, the one dumped by the DrawingTools::DumpSelections() method. It is an explaination of the selection
  // - Pointer to the selection. The argument in the constructor (false) indicates the 
  //   step sequence is not broken when a cut is not passed. In this way we can save intermediate information for events 
  //   not passing the entire selection

  // Add a simple selection with no branches
  sel().AddSelection("kTrackerTutorial",          "tutorial selection",                   new tutorialSelection(false));

  // Add a more complicated selection with branches
  sel().AddSelection("kTrackerTutorialBranches",  "tutorial selection with branches",     new tutorialBranchesSelection(false));

  // Disable the ones not enabled in parameters file
  if (!ND::params().GetParameterI("tutorialAnalysis.Selections.RunSelectionWithoutBranches"))
    sel().DisableSelection("kTrackerTutorial");

  if (!ND::params().GetParameterI("tutorialAnalysis.Selections.RunSelectionWithBranches"))
    sel().DisableSelection("kTrackerTutorialBranches");
}

//********************************************************************
void tutorialAnalysis::DefineCorrections(){
//********************************************************************

  /* Corrections modify some aspect of the input data (real data or MC). 
     The entire analysis will proceed on the modified data
  */

  // Some corrections are defined in baseAnalysis (have a look at baseAnalysis/vXrY/src/baseAnalysis.cxx)
  baseAnalysis::DefineCorrections();

  // ----- We can add here extra corrections ----------

  // A dummy correction which moves forward the position of all global tracks by some mm, as specified in 
  // an external data file (data/tutorialCorrection.dat)
  corr().AddCorrection("tutorialCorrection", new tutorialCorrection());
}

//********************************************************************
void tutorialAnalysis::DefineSystematics(){
//********************************************************************
  
  /*  Systematics will modify the effective number of events passing the selection and the distribution of the observables. 
      
      There are two kind of systematics:
      - Variations: modify the input data (momentum, PID variables, etc). The selection is run on the modified data such that 
        the result of the selection can be different 
      - Weights: do not modify the input data. The selection is not affected. Simply a weight is added to the event. 
        Since events with different values of a given observable (i.e. momentum ) can have different weights, 
	distributions of that observable may change.
   */


  // Some systematics are defined in baseAnalysis (have a look at baseAnalysis/vXrY/src/baseAnalysis.cxx)
  baseAnalysis::DefineSystematics();

  // ----- We can add here extra systematics ----------

  // An example of variation systematic
  AddVariationSystematic(syst(), kTutorialVariation, new tutorialVariationSystematics());

  // A example of weight systematic
  AddWeightSystematic(   syst(), kTutorialWeight,    new tutorialWeightSystematics());
}

//********************************************************************
void tutorialAnalysis::DefineConfigurations(){
//********************************************************************

  /*  A "configuration" is defined by the systematics that are enabled, the number of toys being run and the random seed 
      used to generate the toys. Each configuration has a micro-tree associated in the output file (with the same name)
  */

  // Some configurations are defined in baseAnalysis (have a look at baseAnalysis/vXrY/src/baseAnalysis.cxx)
  baseAnalysis::DefineConfigurations();

  Int_t ntoys = ND::params().GetParameterI("baseAnalysis.Systematics.NumberOfToys");
  Int_t randomSeed = ND::params().GetParameterI("baseAnalysis.Systematics.RandomSeed");

  // Add a new configuration called "tutorial_syst" with the following properties:
  // - the number of toys defined in the baseAnalysis parameters file
  // - the random seed for toy experiments defined in the baseAnalysis parameters file
  // - This configuration has two systematics kTutorialWeight and kTutorialVariation, defined above

  AddConfiguration(conf(), tutorial_syst, ntoys, randomSeed);
  conf().EnableSystematic(kTutorialWeight,    tutorial_syst);  // Enable the kTutorialWeight in the tutorial_syst configuration
  conf().EnableSystematic(kTutorialVariation, tutorial_syst);  // Enable the kTutorialVariation in the tutorial_syst configuration

  // DIsable the configuration when requested
  if (!ND::params().GetParameterI("tutorialAnalysis.Systematics.EnableTutorialSystematics"))
    conf().DisableConfiguration(tutorial_syst);
}

//********************************************************************
void tutorialAnalysis::DefineMicroTrees(bool addBase){
//********************************************************************

  /*  We call Micro-trees to the standard analysis trees appearing in the output file. 
      There is always a Micro-Tree call "default" which should contain the basic info to understand our selection. 
      The user can add extra Micro-Trees by adding configurations to the analysis (see DefineConfigurations method above).

      Here we give an example of different variables that can be added. Have a look at highlandTools/vXrY/src/OutputManager.hxx
      to see all available methods.
  */

  // -------- Add variables to the analysis tree ----------------------

  // Variables from baseAnalysis (run, event, ...)
  if (addBase) baseAnalysis::DefineMicroTrees(addBase);

  // --- Single variables -------
  AddVarD(  output(), selmu_theta,      "muon candidate reconstructed theta");     // Double 
  AddVarF(  output(), selmu_truemom,    "muon candidate true momentum");           // Float
  AddVarI(  output(), selmu_detectors,  "muon candidate detectors");               // Integer
  AddVarC(  output(), selmu_sense,      "muon candidate sense");                   // Char

  AddVarI(  output(), nLongTPCTracks,   "number of long TPC tracks");               // Integer

  // --- Vector variables -------
  // selmu_ntpcs is the counter, which is added automatically as integer variable  
  AddVarVI(output(), selmu_tpc_det,         "muon candidate TPC number",                               selmu_ntpcs);
  AddVarVI(output(), selmu_tpc_nnodes,      "muon candidate #nodes in each TPC",                       selmu_ntpcs);  // Vector of integers
  AddVarVF(output(), selmu_tpc_mom,         "muon candidate reconstructed momentum in each TPC",       selmu_ntpcs);  // Vector of floats
  AddVarVD(output(), selmu_tpc_truemom,     "muon candidate true momentum in each TPC",                selmu_ntpcs);  // Vector of doubles

  // --- 3D and 4D vector variables (i.e. directions and positions) -------
  AddVar4VF(output(), selmu_pos,    "muon candidate reconstructed position");     // 4-Vector of floats    
  AddVar3VF(output(), selmu_dir,    "muon candidate reconstructed direction");    // 3-Vector of floats     

  // --- Matrix variables -------
  //  AddVarMF(output(), selmu_tpc, "",selmu_necals,-30,4);

  // --- 3D and 4D matrix variables (i.e. directions and positions for constituents) -------
  AddVar4MF(output(), selmu_tpc_pos,    "muon candidate true position in each tpc",  selmu_ntpcs);         
  AddVar3MF(output(), selmu_tpc_dir,    "muon candidate true direction in each tpc", selmu_ntpcs);         


  // Now we had toy-dependent variables
  // There could be many variables that are toy-dependent. We don't need to save all of them as toy variables, 
  // but only the ones we are interested in plotting for different toys. 
  // TOY VARIABLES ARE VERY SPACE CONSUMING SO WE SHOULD MINIMISE ITS NUMBER !!!!

  // --- single toy variables -------
  AddToyVarF(output(), selmu_mom,      "muon candidate reconstructed momentum");

  // --- vector toy variables -------
  AddToyVarVF(output(), selmu_tpc_dedx,    "muon candidate dEdx (CT) in each TPC", NMAXTPCS);
}

//********************************************************************
void tutorialAnalysis::DefineTruthTree(){
//********************************************************************

  /*  The "truth" tree also appears in the output file. It contains all interactions in which we are interested in regardless on whether 
      the selection was passed or not. This is the tree that should be used to compute signal efficiencies
  */

  // Variables from baseAnalysis (run, event, ...)
  baseAnalysis::DefineTruthTree();

  //--- muon variables -------
  AddVarF(  output(), truemu_truemom,      "true muon true momentum");
  AddVarF(  output(), truemu_truecostheta, "true muon true cos(theta)");
  AddVar4VF(output(), truemu_truepos,      "true muon true position");
  AddVar3VF(output(), truemu_truedir,      "true muon true direction");
}

//********************************************************************
bool tutorialAnalysis::Process(){
//********************************************************************

  /*  This is the method that runs the selection (for each toy experiment). It will return true or false depending 
      on whether the entire selection is passed or not. 
      In this particular example we can chose the selection (one branch or several branches) changing the parameters file:
      this will change the enabled selections (see DefineSelections above) and then the effect of ApplySelection

      The calls bellow could be replaced by 

      bool ok1 = sel().GetSelection("kTrackerTutorial")->Apply(*_event, box());
      bool ok2 = sel().GetSelection("kTrackerTutorialBranches")->Apply(*_event, box());

      However the using ApplySelection method has the advantange that it checks whether the selection exists or not 
      and whether the selection is enabled or not. Otherwise we would have to do it our selves here
  */

  bool ok1 = sel().ApplySelection("kTrackerTutorial",        *_event, box());
  bool ok2 = sel().ApplySelection("kTrackerTutorialBranches",*_event, box());

  // Return true if any of the two selections is passed
  return (ok1 || ok2);
}

//********************************************************************
void tutorialAnalysis::FillMicroTrees(bool addBase){
//********************************************************************

  /*  In this method we fill all toy-independent variables (all except the ones added with AddToy...) defined in the method DefineMicroTrees. 
      This method is called once all toys has been run, what means that the value of all variables for the last toy will be saved. This is not a problem 
      for variables that are not expected to change from a toy to another.
  */

  // Variables from baseAnalysis (run, event, ...)
  if (addBase) baseAnalysis::FillMicroTreesBase(addBase);

  // Add a variable that is only in the derived box AnaBoxTutorial. We need to cast the base box to have access 
  output().FillVar(nLongTPCTracks, static_cast<AnaBoxTutorial*>(&box())->nLongTPCTracks);      

  // Muon candidate variables
  if (box().HMNtrack){    

    // Since Detectors is not in BaseDataClasses (in psycheCore) but in DataClasses in (highlandTools) we need to cast 
    // the AnaTrackB pointer in box().HMNtrack to the derived class AnaTrack.
    output().FillVar(selmu_detectors,             static_cast<AnaTrack*>(box().HMNtrack)->Detectors);      

    // Cast it to a double because we have defined this variable as double in DefineMicroTrees 
    // (no need to be a double, it is just to have an example of double variable)
    output().FillVar(selmu_theta,                 (Double_t)acos(box().HMNtrack->DirectionStart[2]));

    // This is unnecessary here (we could use 1 and -1 which are smaller types), but serves as an example
    if (box().HMNtrack->DirectionStart[2]>0)
      output().FillVar(selmu_sense,                 "forward");
    else
      output().FillVar(selmu_sense,                 "backward");

    // This is the way we fill 3 and 4-vectors. Must specify the array size when using FillVectorVarFromArray
    output().FillVectorVarFromArray(selmu_pos,    box().HMNtrack->PositionStart,  4);
    output().FillVectorVarFromArray(selmu_dir,    box().HMNtrack->DirectionStart, 3);

    // Properties of the true particle associated to the muon candidate
    if(  box().HMNtrack->TrueTrack ) {
      output().FillVar(selmu_truemom,             box().HMNtrack->TrueTrack->Momentum);
    }

    // Info in all TPCs 
    for (Int_t subdet = 0; subdet<3; subdet++) {
      if (!SubDetId::GetDetectorUsed(box().HMNtrack->Detector, static_cast<SubDetId::SubDetEnum >(subdet+2))) continue;
      AnaTpcTrack* TPCSegment = static_cast<AnaTpcTrack*>(anaUtils::GetSegmentInDet( *box().HMNtrack,static_cast<SubDetId::SubDetEnum >(subdet+2)));
      if (!TPCSegment) continue;
      output().FillVectorVar(selmu_tpc_det,     subdet);
      output().FillVectorVar(selmu_tpc_mom,     TPCSegment->Momentum);
      output().FillVectorVar(selmu_tpc_nnodes,  TPCSegment->NNodes);

      // Must specify the array size when using FillMatrixVarFromArray
      output().FillMatrixVarFromArray(selmu_tpc_pos,    TPCSegment->PositionStart,  4);
      output().FillMatrixVarFromArray(selmu_tpc_dir,    TPCSegment->DirectionStart, 3);

      if (TPCSegment->TrueTrack){
	// Cast it to a double because we have defined this variable as double in DefineMicroTrees
	// (no need to be a double, it is just to have an example of double variable)
	output().FillVectorVar(selmu_tpc_truemom, (Double_t)TPCSegment->TrueTrack->Momentum);
      }
      
      // increment the value of the counter for all TPC variables (selmu_ntpcs)
      // The same effect will have output().IncrementCounter(selmu_ntpcs);
      output().IncrementCounterForVar(selmu_tpc_det);
    }
  }
  
}

//********************************************************************
void tutorialAnalysis::FillToyVarsInMicroTrees(bool addBase){
//********************************************************************

  /*  In this method we fill all toy-dependent variables (the ones added with AddToy...) defined in the method DefineMicroTrees. 
      This method is called at the end of each toy.

      There could be many variables that are toy-dependent. We don't need to save all of them as toy variables, but only the ones we are interested in plotting 
      for different toys. 

      TOY VARIABLES ARE VERY SPACE CONSUMING SO WE SHOULD MINIMISE ITS NUMBER !!!!
  */

  // Fill the common variables
  if (addBase) baseAnalysis::FillToyVarsInMicroTreesBase(addBase);

  //---- variables specific for this analysis -----
  if (box().HMNtrack){
    // Since we are applying a systematic that varies the momentum we need to save the momentum for each toy
    output().FillToyVar(selmu_mom, box().HMNtrack->Momentum);    

    // Info in all TPCs 
    for (Int_t subdet = 0; subdet<3; subdet++) {
      if (!SubDetId::GetDetectorUsed(box().HMNtrack->Detector, static_cast<SubDetId::SubDetEnum >(subdet+2))) continue;
      AnaTpcTrack* TPCSegment = static_cast<AnaTpcTrack*>(anaUtils::GetSegmentInDet( *box().HMNtrack, static_cast<SubDetId::SubDetEnum >(subdet+2)));
      if (!TPCSegment) continue;
      // In principle we need this variable here when PID systematics are run
      output().FillToyVectorVar(selmu_tpc_dedx,   TPCSegment->dEdxMeas, subdet);
    }
  }
}

//********************************************************************
bool tutorialAnalysis::CheckFillTruthTree(const AnaTrueVertex& vtx){
//********************************************************************

  /* To avoid unecessary events in the "truth" tree in this method we define the condition to include or not a given 
     true vertex in the tree. 
  */

  // In this case we only save numu (NuPDG=14) charged current  (0<ReacCode<30) interactions in the FGD1 FV
  bool numuCC=vtx.ReacCode>0 && vtx.ReacCode<30 && vtx.NuPDG==14;// && vtx.LeptonPDG==13;  
  return (anaUtils::InFiducialVolume(SubDetId::kFGD1, vtx.Position, FVDef::FVdefminFGD1,FVDef::FVdefmaxFGD1) && numuCC);
}

//********************************************************************
bool tutorialAnalysis::CheckFillRooTracker(const AnaTrueVertex& vtx){
//********************************************************************

  /* To avoid unecessary events in the "RooTrackerTree" tree in this method we define the condition to include or not a given 
     true vertex in the tree. 
  */

  // In this case we save all numu CC interactions
  bool numuCC=vtx.ReacCode>0 && vtx.ReacCode<30 && vtx.NuPDG==14;
  return numuCC;
}

//********************************************************************
void tutorialAnalysis::FillTruthTree(const AnaTrueVertex& vtx){
//********************************************************************

  // Fill the common variables defined in baseAnalysis/vXrY/src/baseAnalysis.cxx
  baseAnalysis::FillTruthTreeBase(vtx);
  
  // ---- Fill the extra variables ------
  output().FillVar(truemu_truecostheta,(Float_t)cos(utils::ArrayToTVector3(vtx.LeptonDir).Angle(utils::ArrayToTVector3(vtx.NuDir))));
  output().FillVar(truemu_truemom,vtx.LeptonMom);
  output().FillVectorVarFromArray(truemu_truedir, vtx.LeptonDir,3);
  output().FillVectorVarFromArray(truemu_truepos, vtx.Position, 4);
}

//********************************************************************
void tutorialAnalysis::FillCategories(){
//********************************************************************

  /* This method fills the micro-tree variables related with track categories for color drawing. 
     Those variables are added automatically (no need to explicitely add them in DefineMicroTrees) to the 
     micro-trees, but must be filled by the analyser, provided the event and the relevant track (the lepton candidate, 
     that is the HMN track, in this case). 

     If this method is not implemented, the one from the base class (baseAnalysis::FillCategories()) will be called.      
  */

  // For the muon candidate
  anaUtils::FillCategories(_event, static_cast<AnaTrack*>(box().HMNtrack),"");
}
